import os
from dotenv import load_dotenv, find_dotenv # 导入 find_dotenv 帮助定位
from langchain.prompts.example_selector import MaxMarginalRelevanceExampleSelector
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings
from langchain.prompts import FewShotPromptTemplate, PromptTemplate

# 加载 .env 文件中的环境变量 (增强调试)
load_dotenv(dotenv_path=find_dotenv(usecwd=True), verbose=True, override=True)

# 从环境变量加载 API 密钥和基础 URL
api_key = os.getenv("OPENAI_API_KEY")
base_url = os.getenv("OPENAI_BASE_URL")
model = os.getenv("DEFAULT_MODEL")

examples = [
    {"input":"happy","output":"sad"},
    {"input":"tall","output":"short"},
    {"input":"sunny","output":"gloomy"},
    {"input":"windy","output":"calm"},
    {"input":"高兴","output":"悲伤"}
]

#构造提示词模板
example_prompt = PromptTemplate(
    input_variables=["input","output"],
    template="原词：{input}\n反义：{output}"
)

# MMR 会保证相似性和多样性
# 调用MMR
example_selector = MaxMarginalRelevanceExampleSelector.from_examples(
    #传入示例组
    examples,
    #使用openai的嵌入来做相似性搜索
    OpenAIEmbeddings(openai_api_base=base_url,openai_api_key=api_key,model="text-embedding-ada-002"),
    #设置使用的向量数据库是什么
    FAISS,
    #结果条数
    k=2,
)

#使用小样本提示词模版来实现动态示例的调用
mmr_prompt = FewShotPromptTemplate(
    example_selector=example_selector,
    example_prompt=example_prompt,
    prefix="给出每个输入词的反义词",
    suffix="原词：{adjective}\n反义：",
    input_variables=["adjective"]
)

#当我们输入一个描述情绪的词语的时候，应该选择同样是描述情绪的一对示例组来填充提示词模版
print(mmr_prompt.format(adjective="难过"))
